import sys
import cv2
import numpy as np
from PyQt5.QtWidgets import (QApplication, QMainWindow, QTabWidget, QWidget, QVBoxLayout,
                             QHBoxLayout, QLabel, QPushButton, QFileDialog, QSlider,
                             QSpinBox, QDoubleSpinBox, QComboBox, QGroupBox, QGridLayout,
                             QMessageBox, QSplitter, QAction, QMenuBar, QToolBar, QStatusBar)
from PyQt5.QtGui import QImage, QPixmap, QFont, QIcon
from PyQt5.QtCore import Qt, QThread, pyqtSignal, QTimer
import os
from datetime import datetime
import json


class ImageProcessor:
    def __init__(self):
        self.original_image = None
        self.processed_image = None
        self.file_path = None

    def load_image(self, file_path):
        self.file_path = file_path
        self.original_image = cv2.imread(file_path)
        if self.original_image is None:
            return False
        self.processed_image = self.original_image.copy()
        return True

    def get_current_image(self):
        return self.processed_image if self.processed_image is not None else self.original_image

    def save_image(self, file_path=None):
        if file_path:
            self.file_path = file_path
        if self.processed_image is not None and self.file_path:
            cv2.imwrite(self.file_path, self.processed_image)
            return True
        return False

    def rotate(self, angle):
        if self.processed_image is not None:
            rows, cols = self.processed_image.shape[:2]
            M = cv2.getRotationMatrix2D((cols / 2, rows / 2), angle, 1)
            self.processed_image = cv2.warpAffine(self.processed_image, M, (cols, rows))
            return True
        return False

    def resize(self, scale):
        if self.processed_image is not None:
            width = int(self.processed_image.shape[1] * scale)
            height = int(self.processed_image.shape[0] * scale)
            self.processed_image = cv2.resize(self.processed_image, (width, height), interpolation=cv2.INTER_AREA)
            return True
        return False

    def crop(self, x1, y1, x2, y2):
        if self.processed_image is not None:
            self.processed_image = self.processed_image[y1:y2, x1:x2]
            return True
        return False

    def adjust_brightness_contrast(self, brightness=0, contrast=1.0):
        if self.processed_image is not None:
            self.processed_image = cv2.convertScaleAbs(self.processed_image, alpha=contrast, beta=brightness)
            return True
        return False

    def smooth(self, kernel_size=3):
        if self.processed_image is not None:
            self.processed_image = cv2.GaussianBlur(self.processed_image, (kernel_size, kernel_size), 0)
            return True
        return False

    def detect_shapes(self):
        if self.processed_image is not None:
            img = self.processed_image.copy()
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            edges = cv2.Canny(gray, 50, 150)
            contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

            for cnt in contours:
                area = cv2.contourArea(cnt)
                if area > 100:  # 过滤小轮廓
                    approx = cv2.approxPolyDP(cnt, 0.02 * cv2.arcLength(cnt, True), True)
                    if len(approx) == 3:
                        shape = "Triangle"
                    elif len(approx) == 4:
                        shape = "Rectangle"
                    elif len(approx) >= 5:
                        shape = "Polygon"
                    else:
                        shape = "Circle"
                    cv2.drawContours(img, [cnt], -1, (0, 255, 0), 2)
                    cv2.putText(img, shape, (approx[0][0][0], approx[0][0][1]),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
            self.processed_image = img
            return True
        return False

    def face_detection(self):
        if self.processed_image is not None:
            img = self.processed_image.copy()
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

            # 加载人脸检测器
            face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
            faces = face_cascade.detectMultiScale(gray, 1.3, 5)

            for (x, y, w, h) in faces:
                cv2.rectangle(img, (x, y), (x + w, y + h), (255, 0, 0), 2)
                roi_gray = gray[y:y + h, x:x + w]
                roi_color = img[y:y + h, x:x + w]

            self.processed_image = img
            return len(faces) > 0
        return False


class VideoProcessor(QThread):
    frame_ready = pyqtSignal(np.ndarray)
    error_occurred = pyqtSignal(str)
    finished = pyqtSignal()

    def __init__(self, camera_index=0, file_path=None):
        super().__init__()
        self.camera_index = camera_index
        self.file_path = file_path
        self.cap = None
        self.running = False
        self.processing_function = None

    def run(self):
        try:
            if self.file_path:
                self.cap = cv2.VideoCapture(self.file_path)
            else:
                self.cap = cv2.VideoCapture(self.camera_index)

            if not self.cap.isOpened():
                self.error_occurred.emit("无法打开视频源")
                return

            self.running = True
            while self.running:
                ret, frame = self.cap.read()
                if not ret:
                    self.error_occurred.emit("无法读取帧")
                    break

                if self.processing_function:
                    frame = self.processing_function(frame)

                self.frame_ready.emit(frame)
                self.msleep(30)  # 控制帧率

        except Exception as e:
            self.error_occurred.emit(str(e))
        finally:
            self.cleanup()
            self.finished.emit()

    def set_processing_function(self, function):
        self.processing_function = function

    def stop(self):
        self.running = False

    def cleanup(self):
        if self.cap and self.cap.isOpened():
            self.cap.release()


class ImageDisplayWidget(QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.initUI()

    def initUI(self):
        self.main_layout = QVBoxLayout(self)

        self.image_label = QLabel(self)
        self.image_label.setAlignment(Qt.AlignCenter)
        self.image_label.setText("请加载图像")
        self.image_label.setMinimumSize(400, 300)
        self.image_label.setStyleSheet("border: 1px solid #ccc;")

        self.main_layout.addWidget(self.image_label)

    def display_image(self, image):
        if image is None:
            self.image_label.setText("请加载图像")
            return

        if len(image.shape) == 3:  # 彩色图像
            height, width, channel = image.shape
            bytesPerLine = 3 * width
            qImg = QImage(image.data, width, height, bytesPerLine, QImage.Format_RGB888).rgbSwapped()
        else:  # 灰度图像
            height, width = image.shape
            qImg = QImage(image.data, width, height, QImage.Format_Grayscale8)

        pixmap = QPixmap.fromImage(qImg)
        self.image_label.setPixmap(pixmap.scaled(
            self.image_label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))

    def resizeEvent(self, event):
        super().resizeEvent(event)
        if not self.image_label.pixmap() or self.image_label.pixmap().isNull():
            return
        self.image_label.setPixmap(self.image_label.pixmap().scaled(
            self.image_label.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))


class ImageProcessingTab(QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.processor = ImageProcessor()
        self.initUI()

    def initUI(self):
        # 创建主布局
        main_layout = QHBoxLayout(self)

        # 左侧控制面板
        control_panel = QVBoxLayout()

        # 文件操作
        file_group = QGroupBox("文件操作")
        file_layout = QVBoxLayout(file_group)

        self.load_btn = QPushButton("加载图像")
        self.load_btn.clicked.connect(self.load_image)

        self.save_btn = QPushButton("保存图像")
        self.save_btn.clicked.connect(self.save_image)
        self.save_btn.setEnabled(False)

        file_layout.addWidget(self.load_btn)
        file_layout.addWidget(self.save_btn)

        # 图像处理操作
        process_group = QGroupBox("图像处理")
        process_layout = QVBoxLayout(process_group)

        # 旋转控制
        rotate_layout = QHBoxLayout()
        self.rotate_label = QLabel("旋转角度:")
        self.rotate_spin = QSpinBox()
        self.rotate_spin.setRange(-180, 180)
        self.rotate_spin.setValue(0)
        self.rotate_btn = QPushButton("旋转")
        self.rotate_btn.clicked.connect(self.rotate_image)

        rotate_layout.addWidget(self.rotate_label)
        rotate_layout.addWidget(self.rotate_spin)
        rotate_layout.addWidget(self.rotate_btn)

        # 缩放控制
        resize_layout = QHBoxLayout()
        self.resize_label = QLabel("缩放比例:")
        self.resize_spin = QDoubleSpinBox()
        self.resize_spin.setRange(0.1, 5.0)
        self.resize_spin.setValue(1.0)
        self.resize_spin.setSingleStep(0.1)
        self.resize_btn = QPushButton("缩放")
        self.resize_btn.clicked.connect(self.resize_image)

        resize_layout.addWidget(self.resize_label)
        resize_layout.addWidget(self.resize_spin)
        resize_layout.addWidget(self.resize_btn)

        # 亮度对比度控制
        brightness_layout = QHBoxLayout()
        self.brightness_label = QLabel("亮度:")
        self.brightness_slider = QSlider(Qt.Horizontal)
        self.brightness_slider.setRange(-100, 100)
        self.brightness_slider.setValue(0)
        self.brightness_btn = QPushButton("应用")
        self.brightness_btn.clicked.connect(self.adjust_brightness_contrast)

        brightness_layout.addWidget(self.brightness_label)
        brightness_layout.addWidget(self.brightness_slider)
        brightness_layout.addWidget(self.brightness_btn)

        contrast_layout = QHBoxLayout()
        self.contrast_label = QLabel("对比度:")
        self.contrast_slider = QSlider(Qt.Horizontal)
        self.contrast_slider.setRange(50, 200)
        self.contrast_slider.setValue(100)
        self.contrast_btn = QPushButton("应用")
        self.contrast_btn.clicked.connect(self.adjust_brightness_contrast)

        contrast_layout.addWidget(self.contrast_label)
        contrast_layout.addWidget(self.contrast_slider)
        contrast_layout.addWidget(self.contrast_btn)

        # 平滑处理
        smooth_layout = QHBoxLayout()
        self.smooth_label = QLabel("平滑核大小:")
        self.smooth_spin = QSpinBox()
        self.smooth_spin.setRange(1, 15)
        self.smooth_spin.setValue(3)
        self.smooth_spin.setSingleStep(2)  # 必须是奇数
        self.smooth_btn = QPushButton("平滑")
        self.smooth_btn.clicked.connect(self.smooth_image)

        smooth_layout.addWidget(self.smooth_label)
        smooth_layout.addWidget(self.smooth_spin)
        smooth_layout.addWidget(self.smooth_btn)

        # 形状识别
        self.shape_btn = QPushButton("形状识别")
        self.shape_btn.clicked.connect(self.detect_shapes)

        # 人脸识别
        self.face_btn = QPushButton("人脸识别")
        self.face_btn.clicked.connect(self.detect_faces)

        # 添加所有处理控件到布局
        process_layout.addLayout(rotate_layout)
        process_layout.addLayout(resize_layout)
        process_layout.addLayout(brightness_layout)
        process_layout.addLayout(contrast_layout)
        process_layout.addLayout(smooth_layout)
        process_layout.addWidget(self.shape_btn)
        process_layout.addWidget(self.face_btn)

        # 添加所有组到控制面板
        control_panel.addWidget(file_group)
        control_panel.addWidget(process_group)
        control_panel.addStretch(1)

        # 右侧图像显示区域
        self.display_widget = ImageDisplayWidget()

        # 添加到主布局
        main_layout.addLayout(control_panel, 1)
        main_layout.addWidget(self.display_widget, 3)

    def load_image(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择图像", "", "图像文件 (*.png *.jpg *.jpeg *.bmp)")
        if file_path:
            if self.processor.load_image(file_path):
                self.display_widget.display_image(self.processor.get_current_image())
                self.save_btn.setEnabled(True)
                self.setWindowTitle(f"计算机视觉开发平台 - {os.path.basename(file_path)}")
            else:
                QMessageBox.warning(self, "错误", "无法加载图像")

    def save_image(self):
        file_path, _ = QFileDialog.getSaveFileName(
            self, "保存图像", "", "图像文件 (*.png *.jpg *.jpeg *.bmp)")
        if file_path:
            if self.processor.save_image(file_path):
                QMessageBox.information(self, "成功", "图像保存成功")
            else:
                QMessageBox.warning(self, "错误", "无法保存图像")

    def rotate_image(self):
        angle = self.rotate_spin.value()
        if self.processor.rotate(angle):
            self.display_widget.display_image(self.processor.get_current_image())

    def resize_image(self):
        scale = self.resize_spin.value()
        if self.processor.resize(scale):
            self.display_widget.display_image(self.processor.get_current_image())

    def adjust_brightness_contrast(self):
        brightness = self.brightness_slider.value()
        contrast = self.contrast_slider.value() / 100.0
        if self.processor.adjust_brightness_contrast(brightness, contrast):
            self.display_widget.display_image(self.processor.get_current_image())

    def smooth_image(self):
        kernel_size = self.smooth_spin.value()
        if self.processor.smooth(kernel_size):
            self.display_widget.display_image(self.processor.get_current_image())

    def detect_shapes(self):
        if self.processor.detect_shapes():
            self.display_widget.display_image(self.processor.get_current_image())
            QMessageBox.information(self, "完成", "形状识别完成")

    def detect_faces(self):
        if self.processor.face_detection():
            self.display_widget.display_image(self.processor.get_current_image())
            QMessageBox.information(self, "完成", "人脸识别完成")


class VideoProcessingTab(QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.processor = None
        self.is_playing = False
        self.initUI()

    def initUI(self):
        # 创建主布局
        main_layout = QVBoxLayout(self)

        # 控制栏
        control_layout = QHBoxLayout()

        self.open_file_btn = QPushButton("打开视频文件")
        self.open_file_btn.clicked.connect(self.open_video_file)

        self.open_camera_btn = QPushButton("打开摄像头")
        self.open_camera_btn.clicked.connect(self.open_camera)

        self.play_pause_btn = QPushButton("播放")
        self.play_pause_btn.setEnabled(False)
        self.play_pause_btn.clicked.connect(self.toggle_play_pause)

        self.stop_btn = QPushButton("停止")
        self.stop_btn.setEnabled(False)
        self.stop_btn.clicked.connect(self.stop_video)

        self.face_detection_btn = QPushButton("人脸检测")
        self.face_detection_btn.setEnabled(False)
        self.face_detection_btn.clicked.connect(self.toggle_face_detection)

        self.shape_detection_btn = QPushButton("形状检测")
        self.shape_detection_btn.setEnabled(False)
        self.shape_detection_btn.clicked.connect(self.toggle_shape_detection)

        control_layout.addWidget(self.open_file_btn)
        control_layout.addWidget(self.open_camera_btn)
        control_layout.addWidget(self.play_pause_btn)
        control_layout.addWidget(self.stop_btn)
        control_layout.addWidget(self.face_detection_btn)
        control_layout.addWidget(self.shape_detection_btn)

        # 视频显示区域
        self.display_widget = ImageDisplayWidget()

        # 添加到主布局
        main_layout.addLayout(control_layout)
        main_layout.addWidget(self.display_widget)

        self.current_processing_function = None
        self.face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')

    def open_video_file(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择视频文件", "", "视频文件 (*.mp4 *.avi *.mov *.mkv)")
        if file_path:
            self.start_video_processor(file_path=file_path)

    def open_camera(self):
        self.start_video_processor(camera_index=0)

    def start_video_processor(self, camera_index=None, file_path=None):
        self.stop_video()

        self.processor = VideoProcessor(camera_index, file_path)
        self.processor.frame_ready.connect(self.display_widget.display_image)
        self.processor.error_occurred.connect(self.handle_error)
        self.processor.finished.connect(self.video_finished)

        self.play_pause_btn.setEnabled(True)
        self.stop_btn.setEnabled(True)
        self.face_detection_btn.setEnabled(True)
        self.shape_detection_btn.setEnabled(True)

        self.is_playing = True
        self.play_pause_btn.setText("暂停")
        self.processor.start()

    def toggle_play_pause(self):
        if self.processor:
            if self.is_playing:
                self.processor.pause()
                self.play_pause_btn.setText("播放")
            else:
                self.processor.resume()
                self.play_pause_btn.setText("暂停")
            self.is_playing = not self.is_playing

    def stop_video(self):
        if self.processor:
            self.processor.stop()
            self.processor.wait()
            self.processor = None
            self.display_widget.display_image(None)

        self.is_playing = False
        self.play_pause_btn.setText("播放")
        self.play_pause_btn.setEnabled(False)
        self.stop_btn.setEnabled(False)
        self.face_detection_btn.setEnabled(False)
        self.shape_detection_btn.setEnabled(False)

    def toggle_face_detection(self):
        if self.processor:
            if self.current_processing_function == self.detect_faces_in_frame:
                self.current_processing_function = None
                self.processor.set_processing_function(None)
                self.face_detection_btn.setText("人脸检测")
            else:
                self.current_processing_function = self.detect_faces_in_frame
                self.processor.set_processing_function(self.current_processing_function)
                self.face_detection_btn.setText("关闭人脸检测")

    def toggle_shape_detection(self):
        if self.processor:
            if self.current_processing_function == self.detect_shapes_in_frame:
                self.current_processing_function = None
                self.processor.set_processing_function(None)
                self.shape_detection_btn.setText("形状检测")
            else:
                self.current_processing_function = self.detect_shapes_in_frame
                self.processor.set_processing_function(self.current_processing_function)
                self.shape_detection_btn.setText("关闭形状检测")

    def detect_faces_in_frame(self, frame):
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = self.face_cascade.detectMultiScale(gray, 1.3, 5)

        for (x, y, w, h) in faces:
            cv2.rectangle(frame, (x, y), (x + w, y + h), (255, 0, 0), 2)

        return frame

    def detect_shapes_in_frame(self, frame):
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        edges = cv2.Canny(gray, 50, 150)
        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        for cnt in contours:
            area = cv2.contourArea(cnt)
            if area > 100:  # 过滤小轮廓
                approx = cv2.approxPolyDP(cnt, 0.02 * cv2.arcLength(cnt, True), True)
                if len(approx) == 3:
                    shape = "Triangle"
                elif len(approx) == 4:
                    shape = "Rectangle"
                elif len(approx) >= 5:
                    shape = "Polygon"
                else:
                    shape = "Circle"
                cv2.drawContours(frame, [cnt], -1, (0, 255, 0), 2)
                cv2.putText(frame, shape, (approx[0][0][0], approx[0][0][1]),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)

        return frame

    def handle_error(self, error_msg):
        QMessageBox.warning(self, "错误", error_msg)
        self.stop_video()

    def video_finished(self):
        self.is_playing = False
        self.play_pause_btn.setText("播放")


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.initUI()

    def initUI(self):
        # 设置窗口标题和大小
        self.setWindowTitle("计算机视觉开发平台")
        self.resize(1000, 700)

        # 创建标签页（先于菜单和工具栏）
        self.tabs = QTabWidget()

        # 添加图像处理标签页
        self.image_tab = ImageProcessingTab()
        self.tabs.addTab(self.image_tab, "图像处理")

        # 添加视频处理标签页
        self.video_tab = VideoProcessingTab()
        self.tabs.addTab(self.video_tab, "视频处理")

        # 创建菜单栏
        self.create_menu_bar()

        # 创建工具栏
        self.create_tool_bar()

        # 创建状态栏
        self.statusBar().showMessage("就绪")

        # 设置中心部件
        self.setCentralWidget(self.tabs)

    def create_menu_bar(self):
        menubar = self.menuBar()

        # 文件菜单
        file_menu = menubar.addMenu("文件")

        # 加载图像动作
        load_image_action = QAction("加载图像", self)
        load_image_action.setShortcut("Ctrl+I")
        load_image_action.triggered.connect(self.image_tab.load_image)
        file_menu.addAction(load_image_action)

        # 保存图像动作
        save_image_action = QAction("保存图像", self)
        save_image_action.setShortcut("Ctrl+S")
        save_image_action.triggered.connect(self.image_tab.save_image)
        file_menu.addAction(save_image_action)

        file_menu.addSeparator()

        # 退出动作
        exit_action = QAction("退出", self)
        exit_action.setShortcut("Ctrl+Q")
        exit_action.triggered.connect(self.close)
        file_menu.addAction(exit_action)

        # 编辑菜单
        edit_menu = menubar.addMenu("编辑")

        # 撤销动作
        undo_action = QAction("撤销", self)
        undo_action.setShortcut("Ctrl+Z")
        edit_menu.addAction(undo_action)

        # 重做动作
        redo_action = QAction("重做", self)
        redo_action.setShortcut("Ctrl+Y")
        edit_menu.addAction(redo_action)

        # 帮助菜单
        help_menu = menubar.addMenu("帮助")

        # 关于动作
        about_action = QAction("关于", self)
        about_action.triggered.connect(self.show_about_dialog)
        help_menu.addAction(about_action)

    def create_tool_bar(self):
        toolbar = QToolBar("工具栏")
        self.addToolBar(toolbar)

        # 加载图像按钮
        load_image_btn = QAction("加载图像", self)
        load_image_btn.triggered.connect(self.image_tab.load_image)
        toolbar.addAction(load_image_btn)

        # 保存图像按钮
        save_image_btn = QAction("保存图像", self)
        save_image_btn.triggered.connect(self.image_tab.save_image)
        toolbar.addAction(save_image_btn)

        toolbar.addSeparator()

        # 撤销按钮
        undo_btn = QAction("撤销", self)
        toolbar.addAction(undo_btn)

        # 重做按钮
        redo_btn = QAction("重做", self)
        toolbar.addAction(redo_btn)

    def show_about_dialog(self):
        QMessageBox.about(self, "关于",
                          "计算机视觉开发平台\n\n"
                          "版本: 1.0\n"
                          "作者: doubao\n"
                          "日期: 2025年\n\n"
                          "本平台用于演示基础计算机视觉功能，包括图像处理和视频分析。")


if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())